#include <buola/cv/cvisualservoing.h>

#include <buola/scene/cimagerenderer.h>
#include <buola/scene/transform/crttransform.h>
#include <buola/mat/svd.h>
#include <buola/mat/io.h>
#include <buola/image/cimagegc.h>
#include <buola/image/algorithm/color.h>

#include <buola/cv/features/cdetector.h>
#include <buola/cv/features/cmatcher.h>

#include <buola/io/cbinarywriter.h>

namespace buola { namespace cvi {

static int sSobel[][5]=
{
    {-1,-2,0,2,1},
    {-4,-8,0,8,4},
    {-6,-12,0,12,6},
    {-4,-8,0,8,4},
    {-1,-2,0,2,1}
};

static int sNeighborsX[]={0,1,0,-1,0,1,1,-1,-1};
static int sNeighborsY[]={0,0,1,0,-1,1,-1,1,-1};

CPoint_i find_nonzero_neighbor(const img::CView_gray32 &pI,const CPoint_i &pP)
{
    for(int i=0;i<9;i++)
    {
        CPoint_i lP(pP.x+sNeighborsX[i],pP.y+sNeighborsY[i]);
        if(lP.x<0||lP.x>=pI.Width()||lP.y<0||lP.y>=pI.Height()) continue;
        if(pI(lP.x,lP.y)==(unsigned int)-1) continue;
        return lP;
    }
    msg_warn() << "nonzero neighbor not found\n";
    return CPoint_i(-1,-1);
}

CPoint_d transform_to_pixel(const mat::CVec4d &pV)
{
    mat::CVec4d lTrans=pV/pV.w();
    return CPoint_d(lTrans.x()*320.0+319.5,-lTrans.y()*240.0+239.5);
}

mat::CVec4d transform_from_pixel(const CPoint_d &pP,double pDepth,const mat::CMat4d &pM,double pRealDepth)
{
    mat::CVec4d lPoint((pP.x-319.5)/320.0,-(pP.y-239.5)/240.0,pDepth,1);
    lPoint*=-pRealDepth;
    return pM*lPoint;
}

CVisualServoing::CVisualServoing()
    :   mRenderer(new scene::CImageRenderer)
    ,   mDeltas{0.001,0.001,0.001,0.005*M_PI,0.005*M_PI,0.005*M_PI}
    ,   mAlphas{0.5,0.5,0.5,0.5,0.5,0.5}
    ,   mBetas{0.5,0.5,0.5,1.0,1.0,1.0}
{
}

CVisualServoing::~CVisualServoing()
{
}

void CVisualServoing::SetCamera(const scene::PCamera &pCamera)
{
    mCamera=pCamera;
    mRenderer->SetCamera(mCamera);
}

void CVisualServoing::SetScene(const scene::PScene &pScene,const scene::PRTTransform &pControlled)
{
    mScene=pScene;
    mControlled=pControlled;
    mRenderer->AddScene(pScene);
    
    UpdateTransforms();
}

void CVisualServoing::AddControlledJoint(const scene::PRevolutionJoint &pJoint)
{
    mControlledJoints.push_back(pJoint);
}

void CVisualServoing::SetTarget(const img::CView_rgb8 &pTargetImage,const img::CView_gray32f &pTargetDepth)
{
    mTargetImage.Recreate(pTargetImage.Size());
    convert_pixels(pTargetImage,mTargetImage);
    mTargetDepth=clone(pTargetDepth);

    mChamfer.reset(new COrientedChamfer(mTargetImage));
}

void CVisualServoing::SetInitialPose(const mat::CVec3d &pPos,const mat::CQuat_d &pRot)
{
    mPosition=pPos;
    mRotation=pRot;
    
    UpdateTransforms();
}

void CVisualServoing::UpdateTransforms()
{
    if(!mControlled) return;
    
    mControlled->SetTranslation(mPosition);
    mControlled->SetRotation(mRotation);
}

void CVisualServoing::Iterate(bool pReally,bool p2D,bool p3D)
{
    mLastRender.Recreate(mTargetImage.Size());
    mLastDepth.Recreate(mTargetImage.Size());
    mTransformIndex.Recreate(mTargetImage.Size());

    mTransforms.clear();
    
//     if(mControlledJoints.empty())
//     {
//         mRenderer->GetImage(mLastRender,mLastDepth);
//     }
//     else
    {
        mRenderer->GetImage(mLastRender,mLastDepth,mTransformIndex,mTransforms);
    }

    mat::CMat_b lAffectedTransforms(mControlledJoints.size(),mTransforms.size());
    //calculate the relevant matrices
    mat::CMat4d lProjection=mCamera->GetProjectionMatrix();
    mat::CMat4d lVPMatrix=lProjection*mCamera->GetLookAtMatrix();
    mat::CMat4d lVPInverse=inverse(lVPMatrix);
    std::vector<mat::CMat4d> lDMatrices;
    
    mLastVPInverse=lVPInverse;
    mLastProjection=lProjection;
    
    if(mControlledJoints.empty())
    {
        //calculate the relevant matrices
        lDMatrices.emplace_back(lVPMatrix*mat::translation_3d(mat::CVec3d(mDeltas[0],0,0)));
        lDMatrices.emplace_back(lVPMatrix*mat::translation_3d(mat::CVec3d(0,mDeltas[1],0)));
        lDMatrices.emplace_back(lVPMatrix*mat::translation_3d(mat::CVec3d(0,0,mDeltas[2])));

        //matrices for adding a differential rotation to a point in world coordinates
        //X=T*R*Xm
        //Xt=T*R'*T^-1*X=T*R'*T^-1*T*R*Xm=T*R'*R*Xm
        lDMatrices.emplace_back(lVPMatrix*mat::rotation_3d_around(mat::CQuat_d(mDeltas[3],mat::CVec3d(1,0,0)),mPosition));
        lDMatrices.emplace_back(lVPMatrix*mat::rotation_3d_around(mat::CQuat_d(mDeltas[4],mat::CVec3d(0,1,0)),mPosition));
        lDMatrices.emplace_back(lVPMatrix*mat::rotation_3d_around(mat::CQuat_d(mDeltas[5],mat::CVec3d(0,0,1)),mPosition));
    }
    else
    {
        for(int i=0;i<mControlledJoints.size();i++)
        {
            for(int j=0;j<mTransforms.size();j++)
            {
                lAffectedTransforms(i,j)=mControlledJoints[i]==mTransforms[j]||
                                        mControlledJoints[i]->IsAncestorOf(mTransforms[j]);
            }
            
            double lNewValue=mControlledJoints[i]->Value()+mDeltas[i];
            
            mat::CMat34d lOldWorldTransform=mControlledJoints[i]->GetWorldTransform();
            mat::CMat34d lNewWorldTransform=mControlledJoints[i]->GetWorldTransformForValue(lNewValue);
            
            lDMatrices.emplace_back(lVPMatrix*lNewWorldTransform*inverse(lOldWorldTransform));
        }
    }

    mFeatures.clear();
    img::CPixel_rgb8 lBack(127,178,229);

    double lTotalPixelDist=0;
    int lTotalPixelCount=0;
    double lTotalDepthDist=0;
    int lTotalDepthPixelCount=0;

    int lDepthFeatures=0;
    
    for(int x=1;x<639;x++)
    {
        for(int y=1;y<479;y++)
        {
            //TODO:find a better way to calculate the derivative:
            // if the distance to a border is greater than one, use that direction as the approach direction, and
            // calculate the derivative with respect to that direction
            // otherwise, use the sobel filter to find the direction
            if(mLastRender(x,y)==lBack) continue;

            if(p3D&&mLastDepth(x,y)!=(uint32_t)-1&&!is_nan(mTargetDepth(x,y)))
            {
                if(x%4==0&&y%4==0)
                {
                    //add it as a feature
                    SFeature lF;
                    double lFDepth=2.0*mLastDepth(x,y)/((double)(unsigned int)-1)-1.0; //from -1 to 1
                    double lNewDepth=-lProjection(2,3)/(lFDepth+lProjection(2,2));
                    double lTargetDepth=-mTargetDepth(x,y);
                    double lTargetFDepth=(-lProjection(2,3)/lTargetDepth)-lProjection(2,2);

                    if(std::abs(lTargetDepth-lNewDepth)>0.03) continue;
                    lF.mRP={x,y};
                    lF.mOri=-1;
                    lF.mPos=transform_from_pixel(CPoint_d(x,y),lFDepth,lVPInverse,lNewDepth);
                    mat::CVec4d lTargetPos=transform_from_pixel(CPoint_d(x,y),lTargetFDepth,lVPInverse,lTargetDepth);
                    lF.mDist=200.0*std::abs(lTargetPos.z()-lF.mPos.z());
                    lTotalDepthDist+=std::abs(lTargetDepth-lNewDepth);
                    lTotalDepthPixelCount++;
                    for(int i=0;i<lDMatrices.size();i++)
                    {
                        if(!mControlledJoints.empty()&&!lAffectedTransforms(i,mTransformIndex(x,y)))
                        {
                            lF.mDD[i]=0;
                            continue;
                        }
                        
                        mat::CMat4d lMatrix=lVPInverse*lDMatrices[i];
                        mat::CVec4d lNewPos=lMatrix*lF.mPos;
                        double lNewDist=200.0*std::abs(lNewPos.z()-lTargetPos.z());
                        lF.mDD[i]=(lNewDist-lF.mDist)/mDeltas[i];
                    }
                    mFeatures.push_back(lF);
                    lDepthFeatures++;
                }
            }

            if(mLastRender(x+1,y)!=lBack&&mLastRender(x,y+1)!=lBack&&mLastRender(x-1,y)!=lBack&&mLastRender(x,y-1)!=lBack) continue;

            int lXSobel=0,lYSobel=0;

            for(int dx=-2;dx<=2;dx++)
            {
                for(int dy=-2;dy<=2;dy++)
                {
                    lXSobel+=sSobel[dy+2][dx+2]*mLastRender(x+dx,y+dy)[0];
                    lYSobel+=sSobel[dx+2][dy+2]*mLastRender(x+dx,y+dy)[0];
                }
            }
            double lAngle=M_PI+atan2(lYSobel,lXSobel);
            int lWhich=(int)(lAngle/(M_PI/8))%8;

            SFeature lF;
            lF.mRP={x,y};
            lF.mOri=lWhich;
            
            double lFDepth=2.0*mLastDepth(x,y)/((double)(unsigned int)-1)-1.0; //from -1 to 1
            double lRealDepth=-lProjection(2,3)/(lFDepth+lProjection(2,2));

            lF.mPos=transform_from_pixel(CPoint_d(x,y),lFDepth,lVPInverse,lRealDepth);

            lF.mClosest=CPoint_d(mChamfer->GetClosest(lF.mRP,lWhich));
            
            lF.mDist=dist(lF.mClosest,CPoint_d(x,y));

            if(lF.mDist>35.0) continue;
            
            if(lF.mDist==0.0)
            {
                //derivatives with respect to pose movements
                for(int i=0;i<lDMatrices.size();i++)
                {
                    if(!mControlledJoints.empty()&&!lAffectedTransforms(i,mTransformIndex(x,y)))
                    {
                        lF.mDD[i]=0;
                        continue;
                    }
                    CPoint_d lNewPos=transform_to_pixel(lDMatrices[i]*lF.mPos);
                    lF.mDD[i]=dist(lF.mClosest,lNewPos)/mDeltas[i];
                }
            }
            else
            {
                //derivatives with respect to pose movements
                for(int i=0;i<lDMatrices.size();i++)
                {
                    if(!mControlledJoints.empty()&&!lAffectedTransforms(i,mTransformIndex(x,y)))
                    {
                        lF.mDD[i]=0;
                        continue;
                    }
//                    double lNewDist=mChamfer->GetDistAt(transform_to_pixel(lDMatrices[i]*lF.mPos),lWhich);
//                    if(lNewDist<0) continue;
//                    lF.mDD[i]=(lNewDist-lF.mDist)/mDeltas[i];

                    CPoint_d lNewPos=transform_to_pixel(lDMatrices[i]*lF.mPos);
                    CPoint_d lDiff=lNewPos-CPoint_d(x,y);
                    CPoint_d lV=CPoint_d(x,y)-lF.mClosest;
                    lV/=norm(lV);
                    lF.mDD[i]=(lDiff.x*lV.x+lDiff.y*lV.y)/mDeltas[i];
                }
            }

            if(p2D)
                mFeatures.push_back(lF);
            
            lTotalPixelDist+=lF.mDist;
            lTotalPixelCount++;
        }
    }
    
//    msg_info() << "total features: " << mFeatures.size() << "\n";
//    msg_info() << "depth feature: " << lDepthFeatures << "\n";
    
    lTotalPixelDist/=lTotalPixelCount;

    mat::CMat_d lJacobian(lDMatrices.size(),mFeatures.size());
    mat::CRow_d lCorrections(mFeatures.size());
    
    for(int i=0;i<mFeatures.size();i++)
    {
        for(int j=0;j<lDMatrices.size();j++)
            lJacobian(j,i)=mFeatures[i].mDD[j];
        lCorrections[i]=-mFeatures[i].mDist;
    }

    if(pReally)
    {
        mat::CMat_d lIJacobian=pseudo_inverse(lJacobian);
        mat::CRow_d lCart=lCorrections*lIJacobian;

        if(mControlledJoints.empty())
        {
             mPosition+=mat::CVec3d(0,0,mAlphas[2]*lCart[2]);
//             mPosition+=mat::CVec3d(mAlphas[0]*lCart[0],mAlphas[1]*lCart[1],mAlphas[2]*lCart[2]);
//             mRotation=mat::CQuat_d(mAlphas[3]*lCart[3],mat::CVec3d(1,0,0))*
//                             mat::CQuat_d(mAlphas[4]*lCart[4],mat::CVec3d(0,1,0))*
//                             mat::CQuat_d(mAlphas[5]*lCart[5],mat::CVec3d(0,0,1))*mRotation;
        }
        else
        {
            for(int i=0;i<mControlledJoints.size();i++)
            {
                mControlledJoints[i]->SetValue(mControlledJoints[i]->Value()+lCart[i]);
                msg_info() << rad2deg(mControlledJoints[i]->Value()) << "\n";
            }
        }
        
        UpdateTransforms();
    }

    lTotalDepthDist/=lTotalDepthPixelCount;
}

void CVisualServoing::GetContoursImage(const img::CMView_rgb8 &pDst)
{
    const img::CImage_gray8 &lEdges=mChamfer->EdgesImage(0);
    
    for(int i=0;i<pDst.Width();i++)
    {
        for(int j=0;j<pDst.Height();j++)
        {
            if(lEdges(i,j)) //non-edge
            {
                pDst(i,j)={0,0,0};
            }
            else //edge
            {
                pDst(i,j)={0,0,255};
            }
        }
    }

    for(int i=0;i<mFeatures.size();i++)
    {
        CPoint_i lP=mFeatures[i].mRP;
        
        pDst(lP.x,lP.y)[0]=255;
    }
}

void CVisualServoing::GetDirectionsImage(const img::CMView_bgra8 &pDst,const mat::CVec3d &pPos,int pSkip)
{
    for(int k=0;k<8;k+=2)
    {
        const img::CImage_gray8 &lEdges=mChamfer->EdgesImage(k);
        for(int i=0;i<pDst.Width();i++)
        {
            for(int j=0;j<pDst.Height();j++)
            {
                if(lEdges(i,j)) //non-edge
                {
                    if(k==0)
                        pDst(i,j)={0,0,0,255};
                }
                else //edge
                {
                    pDst(i,j)={32*(k+1)-1,0,0,255};
                }
            }
        }
    }

    for(int i=0;i<mFeatures.size();i++)
    {
        if(mFeatures[i].mOri<0) continue;
        CPoint_i lP=mFeatures[i].mRP;
        
        pDst(lP.x,lP.y)[2]=32*(mFeatures[i].mOri+1)-1;
    }

    img::CImageGC<img::TFormat_bgra8> lGC(pDst);
    
    lGC.SetLineWidth(1.0);
    lGC.SetSource({0.0,1.0,0.0});

    for(int i=0;i<mFeatures.size();i++)
    {
        if(mFeatures[i].mOri<0) continue;
        if(i%pSkip) continue;
        CPoint_i lFrom=mFeatures[i].mRP;
        lGC.Path().MoveTo(CPoint_d(lFrom));
        lGC.Path().LineTo(mFeatures[i].mClosest);
        lGC.StrokePath();
    }

    CPoint_d lCenter=GetPointInImage(pPos);
    
    if(!lCenter.Undefined())
    {
        lGC.Path().MoveTo(lCenter-CPoint_d(0,5));
        lGC.Path().LineTo(lCenter+CPoint_d(0,5));
        lGC.Path().MoveTo(lCenter-CPoint_d(5,0));
        lGC.Path().LineTo(lCenter+CPoint_d(5,0));
        lGC.SetSource({0.9,0.9,0.4});
        lGC.StrokePath();
    }
}

void CVisualServoing::DrawRenderEdges(const img::CMView_rgb8 &pDst)
{
    for(int i=0;i<mFeatures.size();i++)
    {
        const auto &lF=mFeatures[i];
        
        if(lF.mOri<0)
        {
//            if(lF.mDist>0.015)
//            {
//                pDst(lF.mRP)={0,255,0};
//            }
        }
        else
        {
            for(int i=-1;i<=1;i++)
                for(int j=-1;j<=1;j++)
                    pDst(lF.mRP+CPoint_i(i,j))={0,0,255};
        }
    }
}

CPoint_d CVisualServoing::GetPointInImage(const mat::CVec3d &pPos)
{
    mat::CMat4d lProjection=mCamera->GetProjectionMatrix();
    mat::CMat4d lVPMatrix=lProjection*mCamera->GetLookAtMatrix();

    mat::CVec3d lWorldPos=mat::translation_3d(mPosition)*(mRotation*pPos);

    return transform_to_pixel(lVPMatrix*mat::CVec4d(lWorldPos.x(),lWorldPos.y(),lWorldPos.z(),1));
}

img::CImage_gray8 CVisualServoing::GetMappableMask()
{
    mLastRender.Recreate(mTargetImage.Size());
    img::CImage_gray32 lDepth(mTargetImage.Size());
    img::CImage_gray8 lTransformIndex(mTargetImage.Size());

    std::vector<scene::PTransform> lTransforms;

    mRenderer->GetImage(mLastRender,lDepth,lTransformIndex,lTransforms);

    mat::CMat4d lProjection=mCamera->GetProjectionMatrix();
//    mat::CMat4d lVPMatrix=lProjection*mCamera->GetLookAtMatrix();
//    mat::CMat4d lVPInverse=inverse(lVPMatrix);

    img::CImage_gray8 lMask(mTargetImage.Size(),img::pixel(0));
    for(int x=1;x<639;x++)
    {
        for(int y=1;y<479;y++)
        {
            if(lDepth(x,y)==(uint32_t)-1) continue;
            double lFDepth=2.0*lDepth(x,y)/((double)(unsigned int)-1)-1.0; //from -1 to 1
            double lCenter=-lProjection(2,3)/(lFDepth+lProjection(2,2));
            
            bool lOk=true;
            
            for(int dx=-1;dx<=1;dx++)
            {
                for(int dy=-1;dy<=1;dy++)
                {
                    if(lDepth(x+dx,y+dy)!=(uint32_t)-1&&!is_nan(mTargetDepth(x+dx,y+dy)))
                    {
                        double lF=2.0*lDepth(x+dx,y+dy)/((double)(unsigned int)-1)-1.0; //from -1 to 1
                        double lPoint=-lProjection(2,3)/(lF+lProjection(2,2));
                        if(abs(lPoint-lCenter)<0.02&&abs((double)-mTargetDepth(x+dx,y+dy)-lCenter)<0.02) continue;
                    }
                    lOk=false;
                    break;
                }
            }
            
            if(lOk)
                lMask(x,y)=255;
        }
    }

    CSURFDetector lSURF;
    mat::CMat_f lOut;
    std::vector<CKeyPoint> lKeyPoints=lSURF.Detect(mTargetImage,/*lMask,*/lOut);

    msg_info() << lKeyPoints.size() << " " << lOut.Rows() << "x" << lOut.Cols() << "\n";

    img::CImage_rgb8 lDst(mTargetImage.Size());
    img::convert_pixels(mTargetImage,lDst);

    for(int i=0;i<lKeyPoints.size();i++)
    {
        if(lMask(lKeyPoints[i].Point_i())!=255)
        {
            msg_info() << "blip\n";
            continue;
        }
        lDst(lKeyPoints[i].Point_i())=img::pixel(0,255,0);
    }

    img::save(lDst,io::file("sift.ppm"));
    
    
    {
        img::CImage_rgb8 lSrc2;
        img::load(lSrc2,io::get_home()/"Documents/phd/vvs/set6/frame_20121122T123225.337511.pcd");
        img::CImage_gray8 lImg2(lSrc2.Size());
        img::convert_pixels(lSrc2,lImg2);
        
        mat::CMat_f lOut2;
        
        std::vector<CKeyPoint> lKeyPoints2=lSURF.Detect(lImg2,lOut2);

        msg_info() << lKeyPoints2.size() << " " << lOut2.Rows() << "x" << lOut2.Cols() << "\n";

        CBFMatcher lMatcher(cv::NORM_L2,true);
        std::vector<CMatch> lMatches=lMatcher.Match(lOut,lOut2);
        msg_info() << lMatches.size() << " matches\n";
        
        img::CImage_rgb8 lDst(mTargetImage.Size(),img::pixel(0,0,0));
        img::copy_pixels(mTargetImage,img::channel_view<img::TFormat_gray8>(lDst,0));
        img::copy_pixels(lImg2,img::channel_view<img::TFormat_gray8>(lDst,2));
        img::CImageGC<img::TFormat_rgb8> lGC(lDst);
        lGC.SetSource(img::CColor(0,1,0));
        int lOkCount=0,lKoCount=0;
        std::sort(lMatches.begin(),lMatches.end());
        for(int i=0;i<lMatches.size();i++)
        {
            if(lMatches[i].mDistance>0.2f) continue;
            if(lMask(lKeyPoints[lMatches[i].mIndex1].Point_i())!=255) continue;
            msg_info() << lKeyPoints[lMatches[i].mIndex2].Size() << " " << lMatches[i].mDistance << "\n";
            auto lPoint1=lKeyPoints[lMatches[i].mIndex2].Point();
            auto lPoint2=lKeyPoints2[lMatches[i].mIndex1].Point();
            bool lOk=abs(lPoint1.x-lPoint2.x)<1 && abs(lPoint1.y-lPoint2.y)<1;
            if(lOk)
                lOkCount++;
            else
                lKoCount++;
            lGC.Path().MoveTo(lPoint1);
            lGC.Path().LineTo(lPoint2);
            lGC.StrokePath();
            lGC.Path().Arc(lPoint1,lKeyPoints[lMatches[i].mIndex2].Size()/2);
            lGC.StrokePath();
        }

        msg_info() << lOkCount << " " << lKoCount << "\n";
        img::save(lDst,io::file("siftdiff.ppm"));
    }

    {
        io::CBinaryWriter lDataOut(io::file("sift.mat"));
        lDataOut << lOut;
    }
    
    return lMask;
}

/*namespace cvi*/ } /*namespace buola*/ }
